#!/usr/bin/env python3
import sys
import ldap
import ldap.modlist as modlist
import logging
from logging.handlers import RotatingFileHandler

# 配置参数
LDAP_HOST = 'ldap://35.78.123.123:1389'
ADMIN_DN = "cn=admin,dc=xxxx,dc=com"
ADMIN_PW = 'A0PYcxxxxx'  # 替换为实际密码
BASE_DN = "ou=users,dc=xxxx,dc=com"
USER_FILE = 'user.txt'

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        RotatingFileHandler('ldap_batch_create.log', maxBytes=1*1024*1024, backupCount=3),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('LDAPBatchCreate')

class LDAPUserManager:
    def __init__(self):
        """初始化LDAP连接"""
        try:
            self.conn = ldap.initialize(LDAP_HOST)
            self.conn.set_option(ldap.OPT_REFERRALS, 0)
            self.conn.protocol_version = ldap.VERSION3
            self.conn.simple_bind_s(ADMIN_DN, ADMIN_PW)
            logger.info("LDAP管理员连接成功")
        except Exception as e:
            logger.error(f"LDAP连接失败: {str(e)}")
            raise

    def _is_valid_username(self, username):
        """验证用户名有效性"""
        # 过滤空值、空格、特殊字符
        if not username or not username.strip() or len(username) > 32:
            return False
        return all(c.isalnum() or c in '_-' for c in username)

    def create_user(self, username):
        """创建单个LDAP用户"""
        if not self._is_valid_username(username):
            logger.warning(f"无效用户名: '{username}' (已跳过)")
            return False

        user_dn = f"cn={username},{BASE_DN}"
        user_attrs = {
            'objectClass': [b'top', b'person', b'organizationalPerson', b'inetOrgPerson'],
            'cn': [username.encode('utf-8')],
            'sn': [username.encode('utf-8')],  # 必须字段
            'uid': [username.encode('utf-8')],
            'userPassword': [username.encode('utf-8')],  # 密码=用户名
            'mail': [f"{username}@coinq.com".encode('utf-8')],
            'displayName': [username.encode('utf-8')]  # 可选添加显示名
        }

        try:
            if self._user_exists(username):
                logger.info(f"用户已存在: {username} (已跳过)")
                return True

            ldif = modlist.addModlist(user_attrs)
            self.conn.add_s(user_dn, ldif)
            logger.info(f"成功创建用户: {username} (DN: {user_dn})")
            return True
        except ldap.LDAPError as e:
            logger.error(f"创建用户 {username} 失败: {str(e)}")
            return False

    def _user_exists(self, username):
        """检查用户是否存在"""
        search_filter = f"(cn={username})"
        try:
            result = self.conn.search_s(
                BASE_DN,
                ldap.SCOPE_ONELEVEL,  # 仅搜索当前层级
                search_filter,
                ['cn']
            )
            return bool(result)
        except ldap.LDAPError as e:
            logger.warning(f"检查用户存在时出错: {str(e)}")
            return False

    def batch_create_from_file(self, filename):
        """从文件批量创建用户"""
        success = failed = skipped = 0
        
        try:
            with open(filename, 'r', encoding='utf-8') as f:
                for line in f:
                    username = line.strip()
                    if not username:
                        continue
                    
                    if self.create_user(username):
                        success += 1
                    else:
                        failed += 1
            
            logger.info(
                f"批量创建完成 | 成功: {success}, 失败: {failed}, 跳过: {skipped}\n"
                f"用户邮箱格式: username@coinq.com\n"
                f"初始密码: 与用户名相同"
            )
            return failed == 0
            
        except FileNotFoundError:
            logger.error(f"文件不存在: {filename}")
            return False
        except Exception as e:
            logger.error(f"读取文件时出错: {str(e)}")
            return False

    def __del__(self):
        if hasattr(self, 'conn'):
            self.conn.unbind()

if __name__ == '__main__':
    if len(sys.argv) > 1:
        USER_FILE = sys.argv[1]  # 允许命令行指定文件

    manager = LDAPUserManager()
    
    if manager.batch_create_from_file(USER_FILE):
        print("所有用户处理完成")
        sys.exit(0)
    else:
        print("部分用户创建失败，请检查日志")
        sys.exit(1)